import os
import pickle
import cv2
import matplotlib.pyplot as plt
import numpy as np
import pandas as pd
np.random.seed(1234)
NUM_TRAIN_IMAGES = 393930
NUM_TEST_IMAGES = 99819
NUM_POINTS_IN_IMAGE = 971
NUM_IMAGES_TO_SHOW = 16
NUM_COLS = 4
NUM_ROWS = NUM_IMAGES_TO_SHOW // NUM_COLS
assert NUM_IMAGES_TO_SHOW == NUM_COLS * NUM_ROWS
RANDOM_TEST_INDICES = np.random.choice(NUM_TEST_IMAGES, size=NUM_IMAGES_TO_SHOW, replace=False)
RANDOM_TRAIN_INDICES = np.random.choice(NUM_TRAIN_IMAGES, size=NUM_IMAGES_TO_SHOW, replace=False)
train_landmarks_df = pd.read_csv("../data/train/landmarks.csv", sep="\t")
assert train_landmarks_df.shape == (NUM_TRAIN_IMAGES, NUM_POINTS_IN_IMAGE * 2 + 1)
def draw_landmarks(image, landmarks):
for point in landmarks:
x, y = point.astype(np.int)
# image, center, radius, color, thickness, lineType
cv2.circle(image, (x, y), 1, (128, 0, 128), 1, -1)
return image
def plot_train_data(is_draw_landmarks=True):
plt.figure(figsize=(25, NUM_ROWS * 8))
for i, index in enumerate(RANDOM_TRAIN_INDICES, 1):
picture_name = train_landmarks_df.iloc[index].iloc[0]
image = cv2.imread(f"../data/train/images/{picture_name}")
image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
if is_draw_landmarks:
landmarks = train_landmarks_df.iloc[index].iloc[1:].values.reshape(-1, 2)
image = draw_landmarks(image, landmarks)
plt.subplot(NUM_ROWS, NUM_COLS, i)
plt.imshow(image)
plt.tight_layout()
plt.show()
def read_test_predictions(pickle_name):
with open(pickle_name, "rb") as fp:
results = pickle.load(fp)
image_names = results["image_names"]
landmarks = results["landmarks"]
assert isinstance(image_names, list)
assert len(image_names) == NUM_TEST_IMAGES
assert isinstance(landmarks, np.ndarray)
assert landmarks.shape == (NUM_TEST_IMAGES, NUM_POINTS_IN_IMAGE, 2)
return image_names, landmarks
def plot_test_predictions(pickle_name, is_draw_landmarks=True):
image_names, landmarks = read_test_predictions(pickle_name)
plt.figure(figsize=(25, NUM_ROWS * 8))
for i, index in enumerate(RANDOM_TEST_INDICES, 1):
picture_name = image_names[index]
image = cv2.imread(f"../{picture_name}")
image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
if is_draw_landmarks:
image = draw_landmarks(image, landmarks[index])
plt.subplot(NUM_ROWS, NUM_COLS, i)
plt.imshow(image)
plt.tight_layout()
plt.show()
plot_train_data(is_draw_landmarks=False)
plot_train_data()
pred_pickle_name = "resnet101_batch_64_test_predictions.pkl"
plot_test_predictions(pred_pickle_name, is_draw_landmarks=False)